# MIT License
#
# Copyright (c) <2009> <Fabio Domingues - fnds3000 in gmail.com>
# Copyright (c) <2011> <Gabriel Tremblay - initnull@gmail.com>
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.

import socket, thread, select
from threading import Thread
import threading
import urllib

__version__ = '0.0.1'
BUFLEN = 8192
VERSION = 'Interceptor Proxy/'+__version__
HTTPVER = 'HTTP/1.1'

class ConnectionHandler:
    # Store session headers
    __headers__ = dict()

    def __init__(self, connection, address, timeout):
        self.client = connection
        self.client_buffer = ''
        self.timeout = timeout
        self.method, self.path, self.protocol = self.get_base_header()
        if self.method=='CONNECT':
            self.method_CONNECT()
        elif self.method=='GET':
            self.method_get()
        elif self.method in ('OPTIONS', 'HEAD', 'POST', 'PUT',
                             'DELETE', 'TRACE'):
            self.method_others()

        self.client.close()
        self.target.close()

    def get_base_header(self):
        end = 0
        while 1:
            self.client_buffer += self.client.recv(BUFLEN)
            end = self.client_buffer.find('\n')
            if end!=-1:
                break

        data = (self.client_buffer[:end+1]).split()
        self.client_buffer = self.client_buffer[end+1:]
        return data

    def method_CONNECT(self):
        self._connect_target(self.path)
        self.client.send(HTTPVER+' 200 Connection established\n'+
                         'Proxy-agent: %s\n\n'%VERSION)
        self.client_buffer = ''
        self._read_write()

    # Headers mangling Get
    def method_get(self):
        self.path = self.path[7:]
        i = self.path.find('/')
        host = self.path[:i]
        path = self.path[i:]
        self._connect_target(host)

        # Strip the last \0x0d\0x0a\0x0d\0x0a
        self.client_buffer = self.client_buffer[:len(self.client_buffer)-4]

        # Split header in lines and build a dict out of it.
        incomingHeaders = {}
        for line in self.client_buffer.splitlines():
            keyValues = line.split(":")
            incomingHeaders[keyValues[0].strip()] = keyValues[1].strip()

        self.client_buffer = ''

        # Add Host value to header
        i = host.find(':')
        self.__headers__['Host'] = host[:i]

        # Apply custom headers over received ones
        for key in self.__headers__.keys():
            incomingHeaders[key] = self.__headers__.get(key)

        # Pop erroneous headers
        if incomingHeaders.get('Range'):
            incomingHeaders.pop('Range')
        if incomingHeaders.get('Icy-MetaData'):
            incomingHeaders.pop('Icy-MetaData')

        # Rebuild header string
        for key in incomingHeaders:
            self.client_buffer = self.client_buffer + key + ": " + incomingHeaders.get(key) + "\r\n"

        # Add back the buffer end
        self.client_buffer += "\r\n"

        # Replace urlencode
        path = urllib.unquote(path)

        self.target.send('%s %s %s\n'%(self.method, path, self.protocol)+
                         self.client_buffer)
        self.client_buffer = ''
        self._read_write()

    def method_others(self):
        self.path = self.path[7:]
        i = self.path.find('/')
        host = self.path[:i]
        path = self.path[i:]
        self._connect_target(host)

        self.target.send('%s %s %s\n'%(self.method, path, self.protocol)+
                         self.client_buffer)
        self.client_buffer = ''
        self._read_write()

    def _connect_target(self, host):
        i = host.find(':')
        if i!=-1:
            port = int(host[i+1:])
            host = host[:i]
        else:
            port = 80
        (soc_family, _, _, _, address) = socket.getaddrinfo(host, port)[0]
        self.target = socket.socket(soc_family)
        self.target.connect(address)

    def _read_write(self):
        time_out_max = self.timeout/3
        socs = [self.client, self.target]
        count = 0
        while 1:
            try:
                count += 1
                (recv, _, error) = select.select(socs, [], socs, 3)
                if error:
                    break
                if recv:
                    for in_ in recv:
                        data = in_.recv(BUFLEN)

                        # Detect Set-Cookie: and assign cookies to proxy custom headers
                        setCookiePos = data.find('Set-Cookie')
                        if setCookiePos > -1:
                            cookies = data[setCookiePos + len('Set-Cookie: '):]
                            end = cookies.find(";")
                            cookies = cookies[:end]
                            self.__headers__['Cookie'] = cookies

                        if data.find('MP2T') > -1:
                            data = data.replace('MP2T', 'mp2t')

                        if in_ is self.client:
                            out = self.target
                        else:
                            out = self.client
                        if data:
                            out.send(data)
                            count = 0
                if count == time_out_max:
                    break
            # Socket error (connection reset by peer / 400 / etc..)
            except socket.error:
                print "Socket error"
                break



class InterceptorProxy(Thread):
    __continue__ = True
    __host__ = 'localhost'
    __port__ = 8080
    __ipv6__ = False
    __timeout__ = 60
    __handler__ = ConnectionHandler
    __proxySocket_ = None
    __injectHeaders__ = {}

    def stopProxy(self):
        print "called"
        self.__continue__ = False
        self.__proxySocket__.close()

    # Start the proxy with parameters
    def startProxy(self, host='localhost', port=8080, IPv6=False, timeout=60, handler=ConnectionHandler,
                   injectHeaders = {}):
        self.__continue__ = True
        self.__host__ = host
        self.__port__ = port
        self.__ipv6__ = IPv6
        self.__timeout__ = timeout
        self.__handler__ = handler
        self.__injectHeaders__ = injectHeaders
        self.start()

    # Run proxy
    def run(self):
        if self.__ipv6__:
            socket_type=socket.AF_INET6
        else:
            socket_type=socket.AF_INET

        self.__proxySocket__ = socket.socket(socket_type)
        self.__proxySocket__.bind((self.__host__, self.__port__))
        self.__proxySocket__.listen(0)

        print "Serving"
        while self.__continue__ is True:
            self.__handler__.__headers__ = self.__injectHeaders__
            thread.start_new_thread(self.__handler__, self.__proxySocket__.accept()+(self.__timeout__,))
